#############################################
# SSH私钥权限修复脚本 (Windows)
# 功能：
# 1. 提示用户输入私钥文件路径
# 2. 验证文件是否存在
# 3. 移除继承权限
# 4. 只授予当前用户完全控制权限
# 5. 验证修复结果
#############################################

# 设置控制台编码为 UTF-8
[Console]::OutputEncoding = [System.Text.Encoding]::UTF8

# 颜色输出函数
function Write-ColorOutput {
    param(
        [string]$Message,
        [string]$Color = "White"
    )
    Write-Host $Message -ForegroundColor $Color
}

function Write-Info {
    param([string]$Message)
    Write-ColorOutput "[INFO] $Message" "Green"
}

function Write-Warn {
    param([string]$Message)
    Write-ColorOutput "[WARN] $Message" "Yellow"
}

function Write-Error-Custom {
    param([string]$Message)
    Write-ColorOutput "[ERROR] $Message" "Red"
}

function Write-Step {
    param([string]$Message)
    Write-ColorOutput "[STEP] $Message" "Cyan"
}

# 显示脚本信息
function Show-Banner {
    Write-Host ""
    Write-Host "=========================================" -ForegroundColor Cyan
    Write-ColorOutput "  SSH私钥权限修复工具 (Windows)" "Cyan"
    Write-Host "=========================================" -ForegroundColor Cyan
    Write-Host ""
}

# 获取用户输入的私钥文件路径
function Get-KeyFilePath {
    Write-Step "请输入SSH私钥文件路径"
    Write-Host ""
    Write-ColorOutput "提示:" "Yellow"
    Write-Host "  - 可以输入绝对路径: C:\Users\Leex\.ssh\id_rsa"
    Write-Host "  - 可以输入相对路径: .ssh\id_rsa 或 id_rsa"
    Write-Host "  - 可以使用波浪号: ~/.ssh/id_rsa"
    Write-Host "  - 可以拖拽文件到此窗口"
    Write-Host ""

    $keyPath = Read-Host "请输入私钥文件路径"

    # 处理空输入
    if ([string]::IsNullOrWhiteSpace($keyPath)) {
        Write-Error-Custom "文件路径不能为空"
        return $null
    }

    # 移除引号（拖拽文件时会带引号）
    $keyPath = $keyPath.Trim('"').Trim("'")

    # 处理波浪号 (~) 路径
    if ($keyPath.StartsWith("~")) {
        $keyPath = $keyPath -replace "^~", $env:USERPROFILE
    }

    # 转换 Unix 风格路径分隔符
    $keyPath = $keyPath -replace "/", "\"

    # 处理相对路径
    if (-not [System.IO.Path]::IsPathRooted($keyPath)) {
        # 如果不是绝对路径，尝试在 .ssh 目录中查找
        $sshDir = Join-Path $env:USERPROFILE ".ssh"
        $possiblePath = Join-Path $sshDir $keyPath

        if (Test-Path $possiblePath) {
            $keyPath = $possiblePath
        } else {
            # 尝试当前目录
            $currentPath = Join-Path (Get-Location) $keyPath
            if (Test-Path $currentPath) {
                $keyPath = $currentPath
            } else {
                $keyPath = $possiblePath  # 使用 .ssh 目录路径，后续会检查是否存在
            }
        }
    }

    # 获取绝对路径
    $keyPath = [System.IO.Path]::GetFullPath($keyPath)

    return $keyPath
}

# 验证文件是否存在
function Test-KeyFile {
    param([string]$FilePath)

    if (-not (Test-Path $FilePath)) {
        Write-Error-Custom "文件不存在: $FilePath"
        return $false
    }

    if ((Get-Item $FilePath).PSIsContainer) {
        Write-Error-Custom "这是一个目录，不是文件: $FilePath"
        return $false
    }

    Write-Info "找到文件: $FilePath"
    return $true
}

# 显示当前文件权限
function Show-CurrentPermissions {
    param([string]$FilePath)

    Write-Step "当前文件权限:"
    Write-Host ""

    try {
        $icaclsOutput = icacls $FilePath 2>&1
        $icaclsOutput | ForEach-Object { Write-Host "  $_" }
        Write-Host ""
    } catch {
        Write-Warn "无法读取当前权限: $_"
    }
}

# 修复文件权限
function Fix-KeyPermissions {
    param([string]$FilePath)

    Write-Step "开始修复文件权限..."
    Write-Host ""

    try {
        # 1. 禁用继承并移除所有现有权限
        Write-Info "步骤 1/3: 移除继承权限..."
        $result = icacls $FilePath /inheritance:r 2>&1
        if ($LASTEXITCODE -ne 0) {
            Write-Error-Custom "移除继承权限失败"
            return $false
        }

        # 2. 只给当前用户完全控制权限
        Write-Info "步骤 2/3: 授予当前用户 ($env:USERNAME) 完全控制权限..."
        $result = icacls $FilePath /grant:r "${env:USERNAME}:F" 2>&1
        if ($LASTEXITCODE -ne 0) {
            Write-Error-Custom "授予权限失败"
            return $false
        }

        # 3. 验证权限
        Write-Info "步骤 3/3: 验证权限设置..."
        Write-Host ""

        return $true

    } catch {
        Write-Error-Custom "修复权限时发生错误: $_"
        return $false
    }
}

# 显示修复后的权限
function Show-FixedPermissions {
    param([string]$FilePath)

    Write-Step "修复后的文件权限:"
    Write-Host ""

    try {
        $icaclsOutput = icacls $FilePath 2>&1
        $icaclsOutput | ForEach-Object {
            if ($_ -match $env:USERNAME) {
                Write-ColorOutput "  $_" "Green"
            } else {
                Write-Host "  $_"
            }
        }
        Write-Host ""

        # 检查是否只有当前用户有权限
        $aclOutput = $icaclsOutput | Where-Object { $_ -match ":" }
        $userCount = ($aclOutput | Measure-Object).Count

        if ($userCount -eq 1 -and $aclOutput -match $env:USERNAME) {
            Write-ColorOutput "✓ 权限设置正确：只有当前用户 ($env:USERNAME) 可以访问" "Green"
        } else {
            Write-Warn "警告：检测到多个用户有访问权限，可能需要手动检查"
        }

    } catch {
        Write-Warn "无法验证权限: $_"
    }
}

# 显示成功摘要
function Show-Summary {
    param([string]$FilePath)

    Write-Host ""
    Write-Host "=========================================" -ForegroundColor Green
    Write-ColorOutput "  权限修复完成!" "Green"
    Write-Host "=========================================" -ForegroundColor Green
    Write-Host ""
    Write-Info "处理的文件:"
    Write-Host "  $FilePath"
    Write-Host ""
    Write-Info "现在可以使用此私钥进行SSH连接了"
    Write-Host ""
}

# 主函数
function Main {
    Show-Banner

    # 获取私钥文件路径
    $keyPath = Get-KeyFilePath

    if ($null -eq $keyPath) {
        Write-Host ""
        Write-Error-Custom "脚本终止"
        exit 1
    }

    Write-Host ""

    # 验证文件存在
    if (-not (Test-KeyFile $keyPath)) {
        Write-Host ""
        Write-Error-Custom "脚本终止"
        exit 1
    }

    Write-Host ""

    # 显示当前权限
    Show-CurrentPermissions $keyPath

    # 确认操作
    Write-Warn "即将修改文件权限，此操作将:"
    Write-Host "  1. 移除所有继承的权限"
    Write-Host "  2. 移除其他用户的访问权限"
    Write-Host "  3. 只保留当前用户 ($env:USERNAME) 的完全控制权限"
    Write-Host ""

    $confirm = Read-Host "确认继续? (yes/no) [默认: yes]"
    if ([string]::IsNullOrWhiteSpace($confirm)) {
        $confirm = "yes"
    }

    if ($confirm -ne "yes" -and $confirm -ne "y") {
        Write-Warn "操作已取消"
        exit 0
    }

    Write-Host ""

    # 修复权限
    $success = Fix-KeyPermissions $keyPath

    if (-not $success) {
        Write-Host ""
        Write-Error-Custom "权限修复失败"
        exit 1
    }

    # 显示修复后的权限
    Show-FixedPermissions $keyPath

    # 显示摘要
    Show-Summary $keyPath
}

# 执行主函数
Main
